import os
import zipfile
from io import BytesIO

import requests


def load_dataset(root_dir:str,dataset_name: str):
    # 说明两个参数
    """
    :param root_dir: 下载解压目录
    :param dataset_name: 下载的数据集名称，不包含zip后缀
    :return:
    """
    # 如果root_dir不为空
    dataset_path = os.path.join(root_dir, dataset_name)
    if os.listdir(root_dir):
        print(f"Directory '{root_dir}' is not empty. Please specify a different directory or delete the existing files.")
        return
    if os.path.exists(dataset_path):
        print(f"Dataset '{dataset_name}' already exists. Skipping download.")
        return





    os.makedirs(root_dir, exist_ok=True)

    minio = "https://minio.yiyunheng.cn:13443/"
    url = minio + f"dataset/图像分类/{dataset_name}.zip"

    print(f"Downloading dataset from {url} ...")
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()  # 如果下载失败会抛异常
    except Exception as e:
        print(f"Failed to download dataset: {e}")
        return

    # 直接用BytesIO在内存中打开zip包，避免写入硬盘zip文件
    with zipfile.ZipFile(BytesIO(response.content)) as zf:
        print(f"Extracting files to {root_dir} ...")
        zf.extractall(root_dir)

    print("Dataset download and extraction complete.")